#!/usr/bin/env python3
"""
通用文档信息抽取 - 智能内存优化训练脚本
手动指定GPU配置并采用最优化的内存策略
支持多种文档类型：合同、发票、简历、医疗报告、身份证等
"""

import os
import sys
import yaml
import logging
import subprocess
from pathlib import Path
from typing import Dict, Any
import time

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class UniversalDocumentTrainer:
    """通用文档信息抽取训练器 - 手动指定GPU配置
    
    支持多种文档类型：
    - 合同文档
    - 发票文档
    - 简历文档
    - 医疗报告
    - 身份证
    - 其他结构化文档
    """
    
    def __init__(self, config_path: str = "config/config.yaml", gpu_ids: str = "0"):
        """初始化训练器
        
        Args:
            config_path: 配置文件路径
            gpu_ids: GPU编号，用逗号分隔 (如 "0", "0,1", "1,2,3")
        """
        self.project_root = Path(__file__).parent.parent
        self.config_path = self.project_root / config_path
        self.config = self._load_config()
        
        # GPU配置
        self._configure_gpus(gpu_ids)
        
        self.master_port = 29500
        
        # 内存优化设置
        self._setup_memory_optimization()
        
        logger.info(f"🔧 环境配置完成: GPU={self.gpu_ids}, 进程数={self.num_gpus}")
    
    def _load_config(self) -> Dict[str, Any]:
        """加载配置文件"""
        if not self.config_path.exists():
            raise FileNotFoundError(f"配置文件不存在: {self.config_path}")
        
        with open(self.config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        
        return config
    
    def _configure_gpus(self, gpu_ids: str):
        """配置GPU资源"""
        try:
            import torch
            
            if not torch.cuda.is_available():
                raise RuntimeError("❌ CUDA不可用")
            
            available_gpus = torch.cuda.device_count()
            logger.info(f"📊 检测到 {available_gpus} 张GPU")
            
            # 清理GPU缓存
            torch.cuda.empty_cache()
            
            # 直接使用指定的GPU编号
            self.gpu_ids = gpu_ids.strip()
            
            # 计算GPU数量
            if ',' in self.gpu_ids:
                self.num_gpus = len(self.gpu_ids.split(','))
            else:
                self.num_gpus = 1
            
            # 验证GPU编号有效性
            gpu_list = [int(gpu.strip()) for gpu in self.gpu_ids.split(',')]
            for gpu_id in gpu_list:
                if gpu_id >= available_gpus:
                    raise ValueError(f"GPU {gpu_id} 不存在，系统只有 {available_gpus} 张GPU")
            
            logger.info(f"🚀 使用GPU: {self.gpu_ids} (共{self.num_gpus}张)")
            
        except Exception as e:
            logger.warning(f"⚠️ GPU配置失败: {e}，使用默认配置")
            self.gpu_ids = "0"
            self.num_gpus = 1
    
    def _setup_memory_optimization(self):
        """设置内存优化环境变量"""
        # PyTorch CUDA内存管理 - 极致优化
        os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True,max_split_size_mb:64,roundup_power2_divisions:8'
        
        # 限制图像处理参数以节省内存 - 大幅降低
        os.environ['MAX_PIXELS'] = '100352'    # 从602112大幅降低
        os.environ['MIN_PIXELS'] = '784'       # 降低最小像素
        os.environ['IMAGE_FACTOR'] = '14'      # 降低图像因子
        os.environ['MAX_RATIO'] = '30'         # 大幅降低最大比例
        
        # CUDA优化设置
        os.environ['CUDA_LAUNCH_BLOCKING'] = '0'
        os.environ['TORCH_CUDNN_V8_API_ENABLED'] = '1'
        os.environ['TORCH_USE_CUDA_DSA'] = '1'
        
        # 线程控制 - 最小化
        os.environ['OMP_NUM_THREADS'] = '1'
        os.environ['MKL_NUM_THREADS'] = '1'
        os.environ['NUMBA_NUM_THREADS'] = '1'
        
        # 禁用不必要的功能
        os.environ['TRANSFORMERS_OFFLINE'] = '1'
        os.environ['HF_DATASETS_OFFLINE'] = '1'
        os.environ['TOKENIZERS_PARALLELISM'] = 'false'
        
        # 分布式训练环境变量
        os.environ['CUDA_VISIBLE_DEVICES'] = self.gpu_ids
        os.environ['WORLD_SIZE'] = str(self.num_gpus)
        os.environ['MASTER_PORT'] = str(self.master_port)
        
        logger.info("🚀 内存优化环境变量已设置")
    
    def _get_data_paths(self):
        """获取数据路径"""
        data_dir = self.project_root / "data"
        
        # 优先使用混合数据集
        mixed_train_data = data_dir / "train_data_mixed_train.jsonl"
        if mixed_train_data.exists():
            logger.info(f"🎆 使用混合文档类型训练数据: {mixed_train_data}")
            return str(mixed_train_data)
        
        # 如果混合数据集不存在，使用原始数据
        train_data = data_dir / "train_data_train.jsonl"
        if train_data.exists():
            logger.info(f"📄 使用原始训练数据: {train_data}")
            return str(train_data)
        
        # 如果都不存在，抛出错误
        raise FileNotFoundError(f"训练数据文件不存在: {train_data} 或 {mixed_train_data}")
    
    def _create_output_dirs(self, stage: int):
        """创建输出目录"""
        models_dir = self.project_root / "models" / "universal_document" / f"stage{stage}"
        logs_dir = self.project_root / "logs" / "universal_document" / f"stage{stage}"
        
        models_dir.mkdir(parents=True, exist_ok=True)
        logs_dir.mkdir(parents=True, exist_ok=True)
        
        return str(models_dir), str(logs_dir)
    
    def _validate_data_format(self) -> bool:
        """验证数据文件格式是否正确"""
        try:
            train_data = self._get_data_paths()
            logger.info(f"📄 检查训练数据: {train_data}")
            
            import json
            with open(train_data, 'r', encoding='utf-8') as f:
                first_line = f.readline().strip()
                if first_line:
                    data = json.loads(first_line)
                    if 'conversation' in data:
                        logger.info("✅ 数据格式验证通过")
                        return True
                    else:
                        logger.error(f"❌ 数据格式错误: 缺少'conversation'字段，数据键: {list(data.keys())}")
                        return False
                else:
                    logger.error("❌ 数据文件为空")
                    return False
        except FileNotFoundError as e:
            logger.error(f"❌ 训练数据文件不存在: {e}")
            return False
        except json.JSONDecodeError as e:
            logger.error(f"❌ 数据文件JSON格式错误: {e}")
            return False
        except Exception as e:
            logger.error(f"❌ 数据验证失败: {e}")
            return False
    
    def _check_environment(self) -> bool:
        """检查环境是否就绪"""
        try:
            logger.info("🔍 开始环境检查...")
            
            # 检查Swift是否安装
            try:
                import swift
                logger.info(f"✅ Swift框架已安装: {swift.__version__}")
            except ImportError:
                logger.error("❌ Swift框架未安装，请先安装: pip install ms-swift")
                return False
            
            # 检查GPU可用性
            import torch
            if not torch.cuda.is_available():
                logger.error("❌ CUDA不可用")
                return False
            
            torch.cuda.empty_cache()
            logger.info("✅ CUDA环境正常")
            
            # 验证数据文件格式
            if not self._validate_data_format():
                logger.error("❌ 数据文件验证失败")
                return False
            
            logger.info("✅ 环境检查全部通过")
            return True
            
        except Exception as e:
            logger.error(f"❌ 环境检查失败: {e}")
            return False
    
    def train_stage1(self):
        """阶段1训练: 视觉-语言对齐层"""
        logger.info("🔥 开始阶段1训练: 视觉-语言对齐层")
        logger.info("-" * 50)
        
        # 获取配置
        stage1_config = self.config['training']['stage1']
        model_config = self.config['model']
        
        # 获取数据和输出路径
        train_data = self._get_data_paths()
        output_dir, log_dir = self._create_output_dirs(1)
        
        # 构建训练命令 - 更激进的内存优化
        base_cmd = [
            'swift', 'sft',
            '--model', model_config['base_model'],
            '--model_type', model_config['model_type'],
            '--train_type', 'full',
            '--dataset', train_data,
            
            # 数据类型和精度
            '--torch_dtype', model_config['torch_dtype'],
            '--attn_impl', self.config['training']['common']['attn_impl'],
            '--bf16', 'true',
            
            # 冻结设置
            '--freeze_vit', str(stage1_config['freeze_vit']).lower(),
            '--freeze_llm', str(stage1_config['freeze_llm']).lower(), 
            '--freeze_aligner', str(stage1_config['freeze_aligner']).lower(),
            
            # 训练参数 - 内存优化
            '--num_train_epochs', str(stage1_config['num_train_epochs']),
            '--per_device_train_batch_size', str(stage1_config['per_device_train_batch_size']),
            '--learning_rate', str(stage1_config['learning_rate']),
            '--gradient_accumulation_steps', str(stage1_config['gradient_accumulation_steps']),
            '--warmup_ratio', str(stage1_config['warmup_ratio']),
            
            # 序列长度限制
            '--max_length', str(model_config['max_length']),
            
            # 输出和日志
            '--output_dir', output_dir,
            '--logging_dir', log_dir,
            
            # 评估和保存 - 启用评估以便查看训练进度
            '--eval_strategy', 'steps',  # 使用Swift支持的参数名
            '--eval_steps', '100',  # 每100步评估一次
            '--save_strategy', 'steps',  # 显式设置保存策略
            '--save_steps', '200',  # 每200步保存一次
            '--save_total_limit', str(self.config['training']['common']['save_total_limit']),
            '--logging_steps', '10',  # 每10步记录一次日志
            '--logging_first_step', 'true',  # 记录第一步
            
            # 数据加载 - 内存优化
            '--dataloader_num_workers', str(self.config['training']['common']['dataloader_num_workers']),
            '--dataset_num_proc', str(self.config['training']['common']['dataset_num_proc']),
            
            # 其他设置
            '--report_to', 'tensorboard',
            '--save_safetensors', 'true',
            '--ddp_find_unused_parameters', 'false',
            '--gradient_checkpointing', 'true',
            '--dataloader_pin_memory', 'false',  # 禁用pin memory以节省内存
            '--remove_unused_columns', 'false',
            
            # 内存优化特殊参数
            '--dataloader_drop_last', 'true',  # 丢弃最后一个批次
            '--fp16_full_eval', 'false',  # 禁用全精度评估
        ]
        
        # 根据GPU数量选择训练方式
        if self.num_gpus == 1:
            # 单GPU训练
            cmd = base_cmd
        else:
            # 多GPU分布式训练 - 使用torchrun
            cmd = [
                'torchrun',
                '--nproc_per_node', str(self.num_gpus),
                '--master_port', str(self.master_port),
                '-m', 'swift.cli.sft'
            ] + base_cmd[2:]  # 去掉 'swift', 'sft'
        
        # 记录命令
        cmd_str = ' '.join(cmd)
        logger.info(f"📋 训练命令: {cmd_str[:100]}...")
        
        try:
            # 执行训练 - 实时显示输出
            logger.info(f"🚀 开始执行训练命令...")
            logger.info(f"📂 工作目录: {self.project_root}")
            logger.info(f"💾 模型输出: {output_dir}")
            logger.info(f"📊 日志输出: {log_dir}")
            
            # 创建日志文件
            log_file = Path(log_dir) / "stage1_training.log"
            log_file.parent.mkdir(parents=True, exist_ok=True)
            
            # 使用实时输出的方式执行训练
            with open(log_file, 'w', encoding='utf-8') as f:
                process = subprocess.Popen(
                    cmd, 
                    cwd=self.project_root,
                    stdout=subprocess.PIPE, 
                    stderr=subprocess.STDOUT,
                    text=True,
                    bufsize=1,
                    universal_newlines=True
                )
                
                # 实时输出训练日志
                logger.info("📋 开始实时输出训练日志:")
                logger.info("-" * 80)
                
                for line in iter(process.stdout.readline, ''):
                    line = line.strip()
                    if line:
                        # 同时输出到控制台和日志文件
                        print(f"[TRAIN] {line}")
                        f.write(f"{line}\n")
                        f.flush()
                        
                        # 特殊关注重要信息
                        if any(keyword in line.lower() for keyword in ['loss', 'epoch', 'step', 'lr', 'gpu', 'error', 'warning']):
                            logger.info(f"📊 {line}")
                
                # 等待进程完成
                return_code = process.wait()
                
                if return_code == 0:
                    logger.info("-" * 80)
                    logger.info("✅ 阶段1训练成功完成!")
                    logger.info(f"📁 模型保存路径: {output_dir}")
                    logger.info(f"📄 完整日志文件: {log_file}")
                    return True
                else:
                    logger.error(f"❌ 阶段1训练失败，退出码: {return_code}")
                    logger.error(f"📄 详细错误日志请查看: {log_file}")
                    return False
            
        except FileNotFoundError as e:
            logger.error(f"❌ 命令未找到: {e}")
            logger.error("请确保已正确安装Swift框架")
            return False
        except KeyboardInterrupt:
            logger.warning("⏹️  训练被用户中断")
            if 'process' in locals():
                process.terminate()
            return False
        except Exception as e:
            logger.error(f"❌ 阶段1训练执行失败: {e}")
            import traceback
            traceback.print_exc()
            return False
    
    def train_stage2(self, stage1_model_path: str = None):
        """阶段2训练: 全模型联合训练"""
        logger.info("🔥 开始阶段2训练: 全模型联合训练")
        logger.info("-" * 50)
        
        # 如果没有指定阶段1模型，使用默认路径
        if stage1_model_path is None:
            stage1_dir = self.project_root / "models" / "universal_document" / "stage1"
            # 查找最新的checkpoint
            checkpoints = list(stage1_dir.glob("**/checkpoint-*"))
            if not checkpoints:
                logger.error("❌ 未找到阶段1的checkpoint，请先运行阶段1训练")
                return False
            stage1_model_path = str(max(checkpoints, key=lambda x: x.stat().st_mtime))
        
        # 获取配置
        stage2_config = self.config['training']['stage2']
        model_config = self.config['model']
        
        # 获取数据和输出路径
        train_data = self._get_data_paths()
        output_dir, log_dir = self._create_output_dirs(2)
        
        # 构建训练命令 - 极致内存优化且兼容Swift
        base_cmd = [
            'swift', 'sft',
            '--model', stage1_model_path,  # 使用阶段1的输出模型
            '--model_type', model_config['model_type'],
            '--train_type', 'lora',  # 使用LoRA而不是full训练
            '--dataset', train_data,
            
            # 数据类型和精度
            '--torch_dtype', model_config['torch_dtype'],
            '--attn_impl', self.config['training']['common']['attn_impl'],
            '--bf16', 'true',
            
            # LoRA配置 - 使用最简化配置
            '--lora_rank', '8',
            '--lora_alpha', '16',
            '--lora_dropout', '0.1',
            
            # 冻结设置（保持部分冻结以节省内存）
            '--freeze_vit', 'true',          # 冻结视觉编码器
            '--freeze_llm', 'false',         # 只训练语言模型
            '--freeze_aligner', 'false',     # 训练对齐层
            
            # 训练参数 - 极致内存优化
            '--num_train_epochs', '1',
            '--per_device_train_batch_size', '1',
            '--learning_rate', '5e-5',
            '--gradient_accumulation_steps', '128',
            '--warmup_ratio', '0.1',
            
            # 序列长度限制
            '--max_length', '2048',
            
            # 输出和日志
            '--output_dir', output_dir,
            '--logging_dir', log_dir,
            
            # 评估和保存 - 启用评估以便监控训练进度
            '--eval_strategy', 'steps',  # 使用Swift支持的参数名
            '--eval_steps', '50',   # 更频繁的评估
            '--save_strategy', 'steps',  # 显式设置保存策略
            '--save_steps', '100',  # 更频繁的保存
            '--save_total_limit', '3',
            '--logging_steps', '5',  # 更频繁的日志
            '--logging_first_step', 'true',
            
            # 数据加载优化
            '--dataloader_num_workers', '0',
            '--dataset_num_proc', '1',
            '--dataloader_pin_memory', 'false',
            '--dataloader_drop_last', 'true',
            
            # 基础设置
            '--save_safetensors', 'true',
            '--gradient_checkpointing', 'true',
            '--remove_unused_columns', 'false',
            '--report_to', 'none',
        ]
        
        # 根据GPU数量选择训练方式
        if self.num_gpus == 1:
            # 单GPU训练
            cmd = base_cmd
        else:
            # 多GPU分布式训练
            cmd = [
                'torchrun',
                '--nproc_per_node', str(self.num_gpus),
                '--master_port', str(self.master_port),
                '-m', 'swift.cli.sft'
            ] + base_cmd[2:]
        
        # 记录命令
        cmd_str = ' '.join(cmd)
        logger.info(f"📋 训练命令: {cmd_str[:100]}...")
        
        try:
            # 执行训练 - 实时显示输出
            logger.info(f"🚀 开始执行阶段2训练命令...")
            logger.info(f"📂 工作目录: {self.project_root}")
            logger.info(f"💾 模型输出: {output_dir}")
            logger.info(f"📊 日志输出: {log_dir}")
            logger.info(f"📝 使用阶段1模型: {stage1_model_path}")
            
            # 创建日志文件
            log_file = Path(log_dir) / "stage2_training.log"
            log_file.parent.mkdir(parents=True, exist_ok=True)
            
            # 使用实时输出的方式执行训练
            with open(log_file, 'w', encoding='utf-8') as f:
                process = subprocess.Popen(
                    cmd, 
                    cwd=self.project_root,
                    stdout=subprocess.PIPE, 
                    stderr=subprocess.STDOUT,
                    text=True,
                    bufsize=1,
                    universal_newlines=True
                )
                
                # 实时输出训练日志
                logger.info("📋 开始实时输出训练日志:")
                logger.info("-" * 80)
                
                for line in iter(process.stdout.readline, ''):
                    line = line.strip()
                    if line:
                        # 同时输出到控制台和日志文件
                        print(f"[TRAIN] {line}")
                        f.write(f"{line}\n")
                        f.flush()
                        
                        # 特殊关注重要信息
                        if any(keyword in line.lower() for keyword in ['loss', 'epoch', 'step', 'lr', 'gpu', 'error', 'warning']):
                            logger.info(f"📊 {line}")
                
                # 等待进程完成
                return_code = process.wait()
                
                if return_code == 0:
                    logger.info("-" * 80)
                    logger.info("✅ 阶段2训练成功完成!")
                    logger.info(f"📁 模型保存路径: {output_dir}")
                    logger.info(f"📄 完整日志文件: {log_file}")
                    return True
                else:
                    logger.error(f"❌ 阶段2训练失败，退出码: {return_code}")
                    logger.error(f"📄 详细错误日志请查看: {log_file}")
                    return False
            
        except FileNotFoundError as e:
            logger.error(f"❌ 命令未找到: {e}")
            logger.error("请确保已正确安装Swift框架")
            return False
        except KeyboardInterrupt:
            logger.warning("⏹️  训练被用户中断")
            if 'process' in locals():
                process.terminate()
            return False
        except Exception as e:
            logger.error(f"❌ 阶段2训练执行失败: {e}")
            import traceback
            traceback.print_exc()
            return False
    
    def run_full_training(self):
        """运行完整的两阶段训练"""
        logger.info("🎯 开始完整的两阶段训练流程")
        logger.info("=" * 60)
        
        # 检查环境
        if not self._check_environment():
            return False
        
        # 阶段1训练
        logger.info("📍 第一阶段: 训练视觉-语言对齐层")
        if not self.train_stage1():
            logger.error("❌ 阶段1训练失败，终止流程")
            return False
        
        # 阶段2训练
        logger.info("📍 第二阶段: 全模型联合训练")
        if not self.train_stage2():
            logger.error("❌ 阶段2训练失败")
            return False
        
        logger.info("🎉 完整训练流程成功完成!")
        return True


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='通用文档信息抽取 - 手动指定GPU训练')
    parser.add_argument('--stage', type=int, choices=[1, 2], 
                       help='指定运行的阶段 (1: 对齐层, 2: 全模型)')
    parser.add_argument('--gpu', type=str, default='0',
                       help='GPU编号，用逗号分隔 (如: 0 或 0,1 或 1,2,3)')
    parser.add_argument('--config', type=str, default='config/config.yaml',
                       help='配置文件路径')
    
    args = parser.parse_args()
    
    try:
        trainer = UniversalDocumentTrainer(args.config, args.gpu)
        
        if args.stage == 1:
            success = trainer.train_stage1()
        elif args.stage == 2:
            success = trainer.train_stage2()
        else:
            # 运行完整流程
            success = trainer.run_full_training()
        
        if success:
            logger.info("🎯 训练任务成功完成!")
            sys.exit(0)
        else:
            logger.error("❌ 训练任务失败!")
            sys.exit(1)
            
    except KeyboardInterrupt:
        logger.info("⏹️  训练被用户中断")
        sys.exit(1)
    except Exception as e:
        logger.error(f"💥 发生未预期的错误: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()